'''
Created on 2019年3月23日

@author: 杜科
'''
#-*- coding: UTF-8 -*-
from __future__ import unicode_literals
import time
import os
import warnings
import numpy as np
import cv2 as cv
import hmmlearn.hmm as hl
from pip._vendor.distlib.compat import raw_input
import pickle


#抑制第三方警告
warnings.filterwarnings(
    'ignore', category=DeprecationWarning)
np.seterr(all='ignore')


def search_objects(directory):
    # 将文件路径格式化，以便适应于各种操作系统
    directory = os.path.normpath(directory)
    # 如果输入路径不是文件则抛出异常
    if not os.path.isdir(directory):
        raise IOError(directory + '不是文件夹')
    objects = {}
    # 遍历文件夹收集jpg文件，放入字典
    for curdir, _, files in os.walk(directory):
        for jpeg in [file for file in files
                     if file.endswith('.jpg')]:
            path = os.path.join(curdir, jpeg)
            # 从文件名中提取图片所属类别
            label = jpeg.split('_')[0]
            # label = path.split(os.path.sep)[-2]
            if label not in objects:
                objects[label] = []
            objects[label].append(path)
    return objects



# 收集图片数据集的标签和灰度值
def label_desc02(objects, flags=None):
    data_x, data_y, data_z = [], [], []
    # 遍历类别标签及对应的文件名
    for label, files in objects.items():
        descs = np.array([])
        for file in files:
            image = cv.imread(file)
            # 训练阶段可不收集图片，不用给定flags值
            if flags:
                # 保存（测试集）中的图片
                data_z.append([])
                data_z[-1].append(image)
            # 将图片转化为灰度图
            gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
            h, w = gray.shape[:2]
            f = 600 / min(h, w)
            gray = cv.resize(gray, None, fx=f, fy=f)
            # 提取灰度图的关键点特征值
            star = cv.xfeatures2d.StarDetector_create()
            keypoints = star.detect(gray)
            sift = cv.xfeatures2d.SIFT_create()
            _, desc = sift.compute(gray, keypoints)
            if len(descs) == 0:
                descs = desc
                # print(descs)
            else:
                descs = np.append(descs, desc, axis=0)
                # print(descs)
            if flags:
                # 测试集按图片进行特征值收集
                data_x.append(descs)
                data_y.append(label)
        if not flags:
            # 训练集按类别进行特征值收集
            data_x.append(descs)
            data_y.append(label)
    return data_x, data_y, data_z

# 可视化识别的图片
def show_pics(test_y, pred_test_y, test_z):
    i = 0
    for label, pred_label, images in zip(
            test_y, pred_test_y, test_z):
        for image in images:
            i += 1
            style = '{} - {} {} {}'.format(
                i, label,
                '==' if label == pred_label
                else '!=', pred_label)
            cv.imshow(style, image)

#训练，返回训练模型            
def train(path):    
    train_path = path
    train_files = search_objects(train_path)
    train_x, train_y, train_z = label_desc02(train_files)       
    return (train_x,train_y)
    
#导入训练模型预测，返回原值与预测值
def test(path,models,train_x,train_y):
    test_path = path
    test_files = search_objects(test_path)      
    test_x, test_y, test_z = label_desc02(test_files, 1)
    models = model_train(train_x, train_y)
    pred_test_y = model_pred(test_x, models)
    return (test_y,pred_test_y)
    
if __name__ == '__main__': 
    path="train"
    train_x,train_y=train(path)
    with open('train_x.txt', 'wb') as f:
        pickle.dump(train_x, f)
    with open('train_y.txt', 'wb') as f:
        pickle.dump(train_y, f)

